package com.lonely.algebra;

import com.google.common.base.Stopwatch;
import com.google.common.collect.Lists;
import com.lonely.base.FoodmartSchema;
import com.lonely.base.ResultSetUtil;
import com.lonely.dialect.Utf8WithMySqlDialect;
import lombok.SneakyThrows;
import org.apache.calcite.adapter.clone.CloneSchema;
import org.apache.calcite.adapter.java.ReflectiveSchema;
import org.apache.calcite.adapter.jdbc.JdbcSchema;
import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.plan.RelTraitDef;
import org.apache.calcite.rel.RelCollation;
import org.apache.calcite.rel.RelCollations;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.JoinRelType;
import org.apache.calcite.rel.rel2sql.RelToSqlConverter;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.schema.Schema;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.sql.*;
import org.apache.calcite.sql.dialect.MysqlSqlDialect;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.parser.SqlParserUtil;
import org.apache.calcite.tools.*;
import org.apache.commons.dbcp2.BasicDataSource;
import org.junit.Test;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * @date: 2022-11-28 13:31
 * @author: ztkj-hzb
 * @description:
 **/
public class AlgebraTest {

    public static Frameworks.ConfigBuilder config() {
        final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
        return Frameworks.newConfigBuilder()
                .parserConfig(SqlParser.Config.DEFAULT)
                .defaultSchema(rootSchema.add("temp", new ReflectiveSchema(new FoodmartSchema())))
                .traitDefs((List<RelTraitDef>) null)
                .programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2));
    }

    public static Frameworks.ConfigBuilder jdbcConfig() {
        final SchemaPlus rootSchema = Frameworks.createRootSchema(true);

        BasicDataSource dataSource = new BasicDataSource();
        dataSource.setDriverClassName("com.mysql.cj.jdbc.Driver");
        dataSource.setUrl("jdbc:mysql://localhost:3306/rpasys");
        dataSource.setUsername("root");
        dataSource.setPassword("duguxiaobiao");

        Schema schema = JdbcSchema.create(rootSchema, "temp", dataSource, null, null);
        SchemaPlus schemaPlus = rootSchema.add("temp", schema);
        //TODO 不知道有什么用，但是去不掉
        SchemaPlus temp1 = rootSchema.add("aa", new CloneSchema(schemaPlus));

        //temp1.add("products_temporal", new ProductsTemporalTable());
        /*temp1.add("orders",
                new OrdersHistoryTable(
                        OrdersStreamTableFactory.getRowList()));*/


        return Frameworks.newConfigBuilder()
                //如果要设置sql参数，在这里设置，比如忽略大小写等
                .parserConfig(SqlParser.Config.DEFAULT)
                .defaultSchema(temp1)
                .traitDefs((List<RelTraitDef>) null)
                .programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2));
    }


    /**
     * value值代数
     */
    @org.junit.Test
    public void testValueAlgebra() {
        final FrameworkConfig config = AlgebraTest.config().build();
        final RelBuilder builder = RelBuilder.create(config);
        builder.values(new String[]{"a", "b"}, 1, true, null, false);
        final RelNode node = builder.build();
        System.out.println(RelOptUtil.toString(node));
    }

    /**
     * 等同： select * from sales_fact_1997;
     */
    @org.junit.Test
    public void testTableScan() {
        final FrameworkConfig config = AlgebraTest.config().build();
        final RelBuilder builder = RelBuilder.create(config);
        final RelNode node = builder
                .scan("sales_fact_1997")
                .build();
        System.out.println(RelOptUtil.toString(node));

        System.out.println("node.explain() = " + node.explain());

    }

    /**
     * project 表示要查询出来的列
     */
    @org.junit.Test
    public void testProject() {
        FrameworkConfig config = AlgebraTest.jdbcConfig().build();
        RelBuilder builder = RelBuilder.create(config);
        RelNode node = builder.scan("emps")
                .project(builder.field("deptNo"), builder.field("name"))
                .build();
        System.out.println(RelOptUtil.toString(node));
    }

    /**
     * 表示where条件
     */
    @Test
    public void testWhere() {
        FrameworkConfig config = AlgebraTest.jdbcConfig().build();
        RelBuilder builder = RelBuilder.create(config);
        RelNode node = builder.scan("emps")
                .filter(builder.equals(builder.field("deptNo"), builder.literal(1)))
                .build();
        System.out.println(RelOptUtil.toString(node));
    }


    /**
     * 集成聚合和having
     */
    @SneakyThrows
    @org.junit.Test
    public void testFilter() {
        FrameworkConfig config = AlgebraTest.jdbcConfig().build();
        RelBuilder builder = RelBuilder.create(config);
        RelNode node = builder.scan("emps")
                .aggregate(builder.groupKey("deptNo"),
                        builder.count(false, "c"),
                        builder.sum(false, "s", builder.field("id")))
                .filter(builder.call(SqlStdOperatorTable.GREATER_THAN, builder.field("c"), builder.literal(10)))
                .build();
        System.out.println(RelOptUtil.toString(node));
        PreparedStatement statement = RelRunners.run(node);
        ResultSet resultSet = statement.executeQuery();
        System.out.println("ResultSetUtil.resultString(resultSet) = " + ResultSetUtil.resultString(resultSet));
    }


    /**
     * 测试如何运行关系代数
     *
     * @throws Exception
     */
    @org.junit.Test
    public void testRun() throws Exception {
        // Equivalent SQL:
        //   SELECT * FROM EMP WHERE DEPTNO = 20
        final RelBuilder builder = RelBuilder.create(jdbcConfig().build());
        RelNode root =
                builder.scan("emps")
                        .filter(
                                builder.equals(builder.field("deptNo"), builder.literal(1)))
                        .build();
        // Note that because the table has been resolved in the RelNode tree
        // we do not need to supply a "schema" as context to the runner.
        try (PreparedStatement preparedStatement = RelRunners.run(root)) {
            ResultSet resultSet = preparedStatement.executeQuery();
            System.out.println("ResultSetUtil.resultString(resultSet) = " + ResultSetUtil.resultString(resultSet));
        }
    }


    /**
     * 测试Join
     */
    @Test
    public void testJoin() throws SQLException {
        final RelBuilder builder = RelBuilder.create(jdbcConfig().build());
        RelNode root =
                builder.scan("emps").as("e")
                        .scan("depts").as("d")
                        //.scan("temp")
                        //左，右，内，全
                        .join(JoinRelType.INNER,
                                builder.equals(builder.field(2, "e", "deptNo"), builder.field(2, "d", "deptNo")))
                        .build();
        // Note that because the table has been resolved in the RelNode tree
        // we do not need to supply a "schema" as context to the runner.
        try (PreparedStatement preparedStatement = RelRunners.run(root)) {
            ResultSet resultSet = preparedStatement.executeQuery();
            System.out.println("ResultSetUtil.resultString(resultSet) = " + ResultSetUtil.resultString(resultSet));
        }
    }


    /**
     * 测试多表join
     */
    @Test
    public void testMultipleTableJoin() throws SQLException {
        RelBuilder relBuilder = RelBuilder.create(jdbcConfig().build());

        RelNode build = relBuilder.scan("emps").as("e")
                .scan("depts").as("d")
                .join(JoinRelType.LEFT, relBuilder.equals(relBuilder.field(2, "e", "deptNo"),
                        relBuilder.field(2, "d", "deptNo")))
                .scan("temp").as("t")
                .join(JoinRelType.LEFT, relBuilder.equals(relBuilder.field(2, "e", "id"),
                        relBuilder.field(2, "t", "销售id")))
                .build();

        String sql = toSql(build);
        System.out.println("sql = " + sql);

        try (PreparedStatement preparedStatement = RelRunners.run(build)) {
            ResultSet resultSet = preparedStatement.executeQuery();
            System.out.println("ResultSetUtil.resultString(resultSet) = " + ResultSetUtil.resultString(resultSet));
        }

    }

    /**
     * 测试union
     */
    @Test
    public void testUnion() throws SQLException {
        RelBuilder relBuilder = RelBuilder.create(jdbcConfig().build());
        RelNode build = relBuilder.scan("emps").as("e")
                .scan("depts").as("d")
                //true : union all  false: union
                .union(true)
                .build();
        System.out.println(toSql(build));

        try (PreparedStatement preparedStatement = RelRunners.run(build)) {
            final ResultSet resultSet = preparedStatement.executeQuery();
            final String resultString = ResultSetUtil.resultString(resultSet);
            System.out.println("resultString = " + resultString);
        }

    }

    /**
     * 测试 查询字段取别名
     */
    @Test
    public void testRename() throws SQLException {
        RelBuilder relBuilder = RelBuilder.create(jdbcConfig().build());
        RelNode build = relBuilder.scan("emps").as("e")
                .scan("depts").as("d")
                .join(JoinRelType.LEFT,
                        relBuilder.equals(relBuilder.field(2, "e", "id"),
                                relBuilder.field(2, "d", "deptNo")))
                .project(relBuilder.field("e", "id"), relBuilder.field("e", "deptNo"),
                        relBuilder.field("d", "id"))
                .rename(Lists.newArrayList("a", "b", "c"))
                .build();
        System.out.println(toSql(build));

        try (PreparedStatement preparedStatement = RelRunners.run(build)) {
            final ResultSet resultSet = preparedStatement.executeQuery();
            final String resultString = ResultSetUtil.resultString(resultSet);
            System.out.println("resultString = " + resultString);
        }

    }


    /**
     * 取交集
     */
    @Test
    public void testIntersect() throws SQLException {
        RelBuilder relBuilder = RelBuilder.create(jdbcConfig().build());
        RelNode build = relBuilder.scan("emps").as("e")
                .scan("depts").as("d")
                //true：不去重 false：去重
                .intersect(true)
                .build();
        System.out.println(toSql(build));

        try (PreparedStatement preparedStatement = RelRunners.run(build)) {
            final ResultSet resultSet = preparedStatement.executeQuery();
            final String resultString = ResultSetUtil.resultString(resultSet);
            System.out.println("resultString = " + resultString);
        }
    }

    /**
     * 取差集
     * @throws SQLException
     */
    @Test
    public void testMinus() throws SQLException {
        RelBuilder relBuilder = RelBuilder.create(jdbcConfig().build());
        RelNode build = relBuilder.scan("emps").as("e")
                .scan("depts").as("d")
                //true：不去重 false：去重
                .minus(true)
                .build();
        System.out.println(toSql(build));

        try (PreparedStatement preparedStatement = RelRunners.run(build)) {
            final ResultSet resultSet = preparedStatement.executeQuery();
            final String resultString = ResultSetUtil.resultString(resultSet);
            System.out.println("resultString = " + resultString);
        }
    }


    /**
     * 变向的笛卡尔集
     * @throws SQLException
     */
    @Test
    public void testCross() throws SQLException {
        Stopwatch stopwatch = Stopwatch.createStarted();
        RelBuilder relBuilder = RelBuilder.create(jdbcConfig().build());
        RelNode build = relBuilder.scan("emps").as("e")
                .scan("depts").as("d")
                .join(JoinRelType.LEFT)
                .build();
        //System.out.println(toSql(build));
        stopwatch.stop();
        final long elapsed = stopwatch.elapsed(TimeUnit.MILLISECONDS);
        System.out.println(elapsed + "ms");

        Stopwatch stopwatch2 = Stopwatch.createStarted();
        try (PreparedStatement preparedStatement = RelRunners.run(build)) {
            final ResultSet resultSet = preparedStatement.executeQuery();
            final String resultString = ResultSetUtil.resultString(resultSet);
            System.out.println("resultString = " + resultString);
        }

        stopwatch2.stop();
        final long elapsed2 = stopwatch2.elapsed(TimeUnit.MILLISECONDS);
        System.out.println(elapsed2 + "ms");
    }





    /**
     * 测试将relNode转换成sql语句
     */
    @Test
    public void testSql() {
        System.setProperty("calcite.default.charset", "UTF-8");
        final RelBuilder builder = RelBuilder.create(jdbcConfig().build());
        RelNode root =
                builder.scan("emps").as("e")
                        .scan("depts").as("d")
                        //.scan("temp")
                        //左，右，内，全
                        .join(JoinRelType.INNER,
                                builder.equals(builder.field(2, "e", "deptNo"), builder.field(2, "d", "deptNo")))
                        .filter(builder.equals(builder.field("e", "name"), builder.literal("独孤")))
                        .build();
        String sql = toSql(root);
        System.out.println("sql = " + sql);
    }


    /**
     * 转换成sql
     *
     * @param root
     * @return
     */
    static String toSql(RelNode root) {
        SqlDialect dialect = Utf8WithMySqlDialect.DEFAULT;
        // 逻辑计划转sql语法树的转换器
        final RelToSqlConverter converter = new RelToSqlConverter(dialect);
        // 把relNode转成sqlNode（sql语法树）
        //final SqlNode sqlNode = converter.visitChild(0, root).asStatement();
        final SqlNode sqlNode = converter.visitRoot(root).asStatement();
        // 把sqlNode转成sql语句
        return sqlNode.toSqlString(dialect).getSql();
    }

}
